import numpy as np

def strassen_multiply(A, B):
    """
    Strassen算法实现矩阵乘法
    :param A: 矩阵A (n x n)
    :param B: 矩阵B (n x n)
    :return: 乘积矩阵C = A × B
    """
    n = A.shape
    
    # 基础情况：当矩阵大小为1x1时直接相乘
    if n == 1:
        return A * B
    
    # 将矩阵分割为四个子矩阵
    mid = n // 2
    A11 = A[:mid, :mid]
    A12 = A[:mid, mid:]
    A21 = A[mid:, :mid]
    A22 = A[mid:, mid:]
    
    B11 = B[:mid, :mid]
    B12 = B[:mid, mid:]
    B21 = B[mid:, :mid]
    B22 = B[mid:, mid:]
    
    # 递归计算7个中间矩阵
    M1 = strassen_multiply(A11 + A22, B11 + B22)
    M2 = strassen_multiply(A21 + A22, B11)
    M3 = strassen_multiply(A11, B12 - B22)
    M4 = strassen_multiply(A22, B21 - B11)
    M5 = strassen_multiply(A11 + A12, B22)
    M6 = strassen_multiply(A21 - A11, B11 + B12)
    M7 = strassen_multiply(A12 - A22, B21 + B22)
    
    # 组合中间结果得到最终矩阵
    C11 = M1 + M4 - M5 + M7
    C12 = M3 + M5
    C21 = M2 + M4
    C22 = M1 - M2 + M3 + M6
    
    # 合并子矩阵
    C = np.vstack((np.hstack((C11, C12)), np.hstack((C21, C22))))
    
    return C

# 示例测试
if __name__ == "__main__":
    A = np.array([[1, 2], [3, 4]])
    B = np.array([[5, 6], [7, 8]])
    
    print("标准乘法结果:")
    print(np.dot(A, B))
    
    print("\nStrassen算法结果:")
    print(strassen_multiply(A, B))